from selenium import webdriver
from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.edge.options import Options as EdgeOptions
from selenium.common.exceptions import WebDriverException

class DriverManager:
    def __init__(self, browser='edge', headless=False):
        self.browser = browser.lower()
        self.headless = headless
        self.driver = None
    
    def init_driver(self):
        """
        初始化WebDriver
        """
        try:
            if self.browser == 'chrome':
                options = ChromeOptions()
                if self.headless:
                    options.add_argument('--headless')
                options.add_argument('--no-sandbox')
                options.add_argument('--disable-dev-shm-usage')
                self.driver = webdriver.Chrome(options=options)
                
            elif self.browser == 'firefox':
                options = FirefoxOptions()
                if self.headless:
                    options.add_argument('--headless')
                self.driver = webdriver.Firefox(options=options)
                
            elif self.browser == 'edge':
                options = EdgeOptions()
                if self.headless:
                    options.add_argument('--headless')
                self.driver = webdriver.Edge(options=options)
                
            else:
                raise ValueError(f"Unsupported browser: {self.browser}")
                
            # 设置隐式等待
            self.driver.implicitly_wait(10)
            return self.driver
        except WebDriverException as e:
            error_msg = f"无法初始化 {self.browser} 浏览器驱动: {str(e)}"
            if "github.com" in str(e):
                error_msg += "\n这通常是由于网络连接问题导致无法下载浏览器驱动。请尝试以下解决方案："
                error_msg += "\n1. 检查网络连接"
                error_msg += "\n2. 配置系统代理（如果使用代理）"
                error_msg += "\n3. 手动下载并配置Firefox驱动 (geckodriver)"
                error_msg += "\n4. 临时使用其他浏览器（如Chrome或Edge）进行测试"
            raise Exception(error_msg) from e
        except Exception as e:
            raise Exception(f"初始化 {self.browser} 浏览器时发生未知错误: {str(e)}") from e
    
    def quit_driver(self):
        """
        关闭WebDriver
        """
        if self.driver:
            self.driver.quit()
            
    def get_driver(self):
        """
        获取WebDriver实例
        """
        if not self.driver:
            self.init_driver()
        return self.driver